import cv2
import os
import numpy as np
from services import db
from datetime import datetime

DATASET_DIR = "images"
MODEL_DIR = "models"
MODEL_PATH = os.path.join(MODEL_DIR, "face_model.yml")
os.makedirs(DATASET_DIR, exist_ok=True)
os.makedirs(MODEL_DIR, exist_ok=True)


# def collect_samples(emp_id, max_samples=30):
#     """采集样本并返回数量"""
#     emp_dir = os.path.join(DATASET_DIR, emp_id)
#     os.makedirs(emp_dir, exist_ok=True)
#     cap = cv2.VideoCapture(0)
#     if not cap.isOpened():
#         return 0
#     face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
#     count = 0
#     while True:
#         ret, frame = cap.read()
#         if not ret:
#             break
#         gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
#         faces = face_cascade.detectMultiScale(gray, 1.3, 5)
#         for (x, y, w, h) in faces:
#             count += 1
#             face_img = gray[y:y + h, x:x + w]
#             cv2.imwrite(os.path.join(emp_dir, f"{count}.jpg"), face_img)
#             cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
#         cv2.imshow("采集样本", frame)
#         if cv2.waitKey(1) & 0xFF == ord('q') or count >= max_samples:
#             break
#     cap.release()
#     cv2.destroyAllWindows()
#     db.update_employee(emp_id, sample_count=count)
#     return count


def save_sample(emp_id, face_img):
    """
    保存一张人脸样本到 datasets/faces/emp_id/ 目录下
    并更新数据库中的 sample_count
    """
    # 创建员工目录
    emp_dir = os.path.join(DATASET_DIR, str(emp_id))
    os.makedirs(emp_dir, exist_ok=True)

    # 文件名用时间戳避免重复
    filename = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + ".jpg"
    filepath = os.path.join(emp_dir, filename)

    # 保存图片
    cv2.imwrite(filepath, face_img)

    # 统计该员工已有样本数
    sample_count = len([f for f in os.listdir(emp_dir) if f.endswith(".jpg")])

    # 更新数据库
    db.update_sample_count(emp_id, sample_count)

    return filepath, sample_count  # 返回路径和当前样本数量


def preview_samples(emp_id, limit=5):
    """返回图片路径列表"""
    emp_dir = os.path.join(DATASET_DIR, emp_id)
    if not os.path.exists(emp_dir):
        return []
    files = sorted(os.listdir(emp_dir))[:limit]
    return [os.path.join(emp_dir, f) for f in files]


def train_and_save_model():
    faces = []
    labels = []
    label_map = {}
    emp_dirs = [d for d in os.listdir(DATASET_DIR) if os.path.isdir(os.path.join(DATASET_DIR, d))]
    if not emp_dirs:
        return False, "没有采集样本"
    label_id = 0
    for emp_id in emp_dirs:
        emp_path = os.path.join(DATASET_DIR, emp_id)
        files = os.listdir(emp_path)
        for file in files:
            img_path = os.path.join(emp_path, file)
            img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
            if img is None:
                continue
            faces.append(img)
            labels.append(label_id)
        label_map[label_id] = emp_id
        label_id += 1
    if not faces:
        return False, "样本无效"
    recognizer = cv2.face.LBPHFaceRecognizer_create()
    recognizer.train(faces, np.array(labels))
    recognizer.write(MODEL_PATH)
    with open(MODEL_PATH + ".labels", "w") as f:
        for k, v in label_map.items():
            f.write(f"{k}:{v}\n")
    return True, "模型训练完成"


def load_model():
    if not os.path.exists(MODEL_PATH):
        return None, {}
    recognizer = cv2.face.LBPHFaceRecognizer_create()
    recognizer.read(MODEL_PATH)
    label_map = {}
    if os.path.exists(MODEL_PATH + ".labels"):
        with open(MODEL_PATH + ".labels", "r") as f:
            for line in f:
                k, v = line.strip().split(":")
                label_map[int(k)] = v
    return recognizer, label_map
